#!/usr/bin/env python3
"""
完整的Sim2Real统一控制器
将humanoid-general-motion-tracking的仿真策略部署到真实Unitree G1机器人

=== 数据读取分类 ===
1. 真实机器人数据读取 (_read_real_robot_state):
   - 数据来源: Unitree G1机器人硬件传感器
   - SDK调用: unitree_sdk2_python DDS通信
   - 传感器: 23个电机编码器 + 9轴IMU传感器
   
2. 模拟数据读取 (_read_simulation_state):
   - 数据来源: 仿真测试模式 - 零状态模拟数据
   - SDK调用: 无 (纯Python数据生成)
   - 用途: 开发测试、算法验证、离线调试

=== 控制指令分类 ===
1. 真实机器人控制 (_send_real_robot_commands):
   - 控制目标: Unitree G1机器人23个电机硬件
   - SDK调用: unitree_sdk2_python DDS通信发送
   - 物理效果: 伺服电机执行PD控制指令
   
2. 模拟控制 (_send_simulation_commands):
   - 控制目标: 仿真测试模式 - 无实际硬件控制
   - SDK调用: 无 (纯Python日志记录)
   - 用途: 算法验证、离线测试、开发调试

核心数据流转:
1. 读取G1机器人23个电机状态 (位置、速度、扭矩、温度)
2. 处理IMU数据 (四元数、角速度、加速度)
3. 策略网络推理生成控制动作
4. PD控制器计算电机扭矩命令
5. 安全检查后发送到机器人
6. 50Hz实时控制循环
"""

import sys, os, time, json
import numpy as np
import torch
from collections import deque
import logging
from datetime import datetime

# Unitree SDK (如果不可用则仿真模式)
try:
    from unitree_sdk2_python.core.channel import ChannelSubscriber, ChannelPublisher
    from unitree_sdk2_python.idl.default import unitree_hg_msg_dds__LowState_, unitree_hg_msg_dds__LowCmd_
    from unitree_sdk2_python.utils.crc import CRC
    UNITREE_AVAILABLE = True
except ImportError:
    print("Unitree SDK unavailable, using simulation mode")
    UNITREE_AVAILABLE = False

from utils.motion_lib import MotionLib

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class G1Sim2RealController:
    """统一的Sim2Real控制器"""
    
    def __init__(self, policy_path, motion_path, interface="lo", domain_id=1, device="cpu"):
        self.interface = interface
        self.domain_id = domain_id 
        self.device = device
        
        # G1机器人参数
        self.num_dofs = 23
        self.control_dt = 0.02  # 50Hz
        self.action_scale = 0.5 if interface == "lo" else 0.3  # 真实机器人更保守
        
        # 默认关节位置
        self.default_dof_pos = np.array([
            -0.2, 0.0, 0.0, 0.4, -0.2, 0.0,  # 左腿
            -0.2, 0.0, 0.0, 0.4, -0.2, 0.0,  # 右腿
            0.0, 0.0, 0.0,                     # 躯干
            0.0, 0.4, 0.0, 1.2,               # 左臂
            0.0, -0.4, 0.0, 1.2,              # 右臂
        ])
        
        # PD控制器参数
        self.stiffness = np.array([100,100,100,150,40,40]*2 + [150,150,150] + [40,40,40,40]*2)
        self.damping = np.array([2,2,2,4,2,2]*2 + [4,4,4] + [5,5,5,5]*2)
        self.torque_limits = np.array([88,139,88,139,50,50]*2 + [88,50,50] + [25,25,25,25]*2)
        
        # 观测历史
        self.history_len = 20
        self.n_proprio = 3 + 2 + 3*self.num_dofs
        self.proprio_history_buf = deque(maxlen=self.history_len)
        for _ in range(self.history_len):
            self.proprio_history_buf.append(np.zeros(self.n_proprio))
            
        # 运动跟踪
        self.tar_obs_steps = torch.tensor([1,5,10,15,20,25,30,35,40,45,50,55,60,65,70,75,80,85,90,95], 
                                         device=self.device, dtype=torch.int)
        
        # 加载策略和动作库
        self.policy_jit = torch.jit.load(policy_path, map_location=self.device)
        self._motion_lib = MotionLib(motion_path, self.device)
        
        # 初始化通信
        self._init_communication()
        
        # 状态变量
        self.last_action = np.zeros(self.num_dofs, dtype=np.float32)
        self.curr_timestep = 0
        self.running = False
        
        # 数据记录
        self.data_log = []
        
        logging.info(f"Controller initialized: interface={interface}, mode={'Real' if UNITREE_AVAILABLE else 'Sim'}")
    
    def _init_communication(self):
        """初始化通信"""
        if UNITREE_AVAILABLE:
            self.state_sub = ChannelSubscriber("rt/lowstate", unitree_hg_msg_dds__LowState_)
            self.cmd_pub = ChannelPublisher("rt/lowcmd", unitree_hg_msg_dds__LowCmd_)
            
            # 初始化命令消息
            self.low_cmd = unitree_hg_msg_dds__LowCmd_()
            self.low_cmd.head[0] = 0xFE
            self.low_cmd.head[1] = 0xEF
            self.low_cmd.level_flag = 0xFF
            
            for i in range(23):
                self.low_cmd.motor_cmd[i].mode = 0x01
                self.low_cmd.motor_cmd[i].q = self.default_dof_pos[i]
                self.low_cmd.motor_cmd[i].kp = self.stiffness[i]
                self.low_cmd.motor_cmd[i].kd = self.damping[i]
    
    def _read_robot_state(self):
        """
        机器人状态数据读取 - 区分真实与模拟数据源
        """
        if not UNITREE_AVAILABLE:
            return self._read_simulation_state()
        else:
            return self._read_real_robot_state()
    
    def _read_simulation_state(self):
        """
        模拟数据读取 - 用于测试和开发
        
        【数据来源】: 仿真测试模式 - 零状态模拟数据
        【SDK调用】: 无 (纯Python数据生成)
        【用途】: 开发测试、算法验证、离线调试
        """
        logging.debug("📊 数据读取模式: 仿真模拟数据")
        
        # 生成模拟的机器人状态数据
        return {
            # 【模拟数据】23个关节的零位置状态
            'joint_pos': np.zeros(23, dtype=np.float32),
            
            # 【模拟数据】23个关节的零速度状态  
            'joint_vel': np.zeros(23, dtype=np.float32),
            
            # 【模拟数据】标准姿态四元数 (无旋转)
            'imu_quat': np.array([0,0,0,1], dtype=np.float32),
            
            # 【模拟数据】零角速度
            'imu_gyro': np.zeros(3, dtype=np.float32),
            
            # 【模拟数据】标准重力加速度 (仅z轴)
            'imu_accel': np.array([0,0,-9.81], dtype=np.float32),
            
            'timestamp': time.time(),
            'data_source': 'simulation_mock',
            'read_mode': 'simulation'
        }
    
    def _read_real_robot_state(self):
        """
        真实机器人数据读取 - 从G1机器人硬件获取
        
        【数据来源】: Unitree G1机器人硬件传感器
        【SDK调用】: unitree_sdk2_python DDS通信
        【传感器详情】:
        - 23个电机: 磁编码器(位置) + 速度估算 + 扭矩传感器 + 温度传感器
        - IMU传感器: 9轴IMU (3轴陀螺仪 + 3轴加速度计 + 3轴磁力计)
        """
        logging.debug("🤖 数据读取模式: 真实机器人硬件")
        
        # 【SDK调用】unitree_sdk2_python.ChannelSubscriber.read()
        # 【数据来源】G1机器人DDS实时状态消息
        state = self.state_sub.read()
        if state is None:
            logging.warning("⚠️ 未接收到机器人状态数据")
            return None
            
        return {
            # 【硬件数据】G1机器人23个电机编码器位置读数
            # 【物理传感器】磁编码器绝对位置测量
            # 【SDK调用】unitree_sdk2_python -> motor_state[i].q
            'joint_pos': np.array([m.q for m in state.motor_state[:23]], dtype=np.float32),
            
            # 【硬件数据】G1机器人23个电机速度反馈
            # 【物理传感器】编码器位置微分计算 + 低通滤波
            # 【SDK调用】unitree_sdk2_python -> motor_state[i].dq  
            'joint_vel': np.array([m.dq for m in state.motor_state[:23]], dtype=np.float32),
            
            # 【硬件数据】G1机器人23个电机实际扭矩
            # 【物理传感器】扭矩传感器测量值 或 电流估算值
            # 【SDK调用】unitree_sdk2_python -> motor_state[i].tau
            'joint_torque': np.array([m.tau for m in state.motor_state[:23]], dtype=np.float32),
            
            # 【硬件数据】G1机器人23个电机温度
            # 【物理传感器】电机驱动器内置温度传感器
            # 【SDK调用】unitree_sdk2_python -> motor_state[i].temperature
            'joint_temperature': np.array([m.temperature for m in state.motor_state[:23]], dtype=np.int32),
            
            # 【硬件数据】G1机器人IMU姿态四元数
            # 【物理传感器】9轴IMU融合算法输出 (陀螺仪+加速度计+磁力计)
            # 【SDK调用】unitree_sdk2_python -> imu_state.quaternion
            # 【坐标系转换】unitree格式[w,x,y,z] -> 标准格式[x,y,z,w]
            'imu_quat': np.array([state.imu_state.quaternion[1], state.imu_state.quaternion[2], 
                                 state.imu_state.quaternion[3], state.imu_state.quaternion[0]], dtype=np.float32),
            
            # 【硬件数据】G1机器人IMU三轴角速度
            # 【物理传感器】3轴MEMS陀螺仪直接测量值
            # 【SDK调用】unitree_sdk2_python -> imu_state.gyroscope
            'imu_gyro': np.array([state.imu_state.gyroscope[0], state.imu_state.gyroscope[1], 
                                 state.imu_state.gyroscope[2]], dtype=np.float32),
            
            # 【硬件数据】G1机器人IMU三轴线性加速度
            # 【物理传感器】3轴MEMS加速度计测量值 (包含重力加速度)
            # 【SDK调用】unitree_sdk2_python -> imu_state.accelerometer
            'imu_accel': np.array([state.imu_state.accelerometer[0], state.imu_state.accelerometer[1], 
                                  state.imu_state.accelerometer[2]], dtype=np.float32),
            
            'timestamp': time.time(),
            'data_source': 'unitree_g1_hardware',
            'read_mode': 'real_robot'
        }
    
    def _get_mimic_obs(self, curr_time_step):
        """获取运动模仿观测
        
        数据来源说明:
        - 运动数据: PKL文件中的预录制人类动作数据
        - SDK调用: 自定义motion_lib.calc_motion_frame()
        """
        num_steps = len(self.tar_obs_steps)
        motion_times = torch.tensor([curr_time_step * self.control_dt], device=self.device).unsqueeze(-1)
        obs_motion_times = self.tar_obs_steps * self.control_dt + motion_times
        obs_motion_times = obs_motion_times.flatten()
        motion_ids = torch.zeros(num_steps, dtype=torch.int, device=self.device)
        
        # 【SDK调用】自定义motion_lib.calc_motion_frame() - 从 PKL动作数据获取轨迹
        # 【数据来源】PKL文件中的人类动作捕捉数据
        root_pos, root_rot, root_vel, root_ang_vel, dof_pos, _ = self._motion_lib.calc_motion_frame(
            motion_ids, obs_motion_times)
        
        # 转换坐标系 - 使用PyTorch工具函数
        # 【SDK调用】utils.torch_utils.quat_rotate_inverse
        from utils.torch_utils import quat_rotate_inverse
        
        def euler_from_quaternion(quat):
            x,y,z,w = quat[:,0], quat[:,1], quat[:,2], quat[:,3]
            roll = torch.atan2(2*(w*x + y*z), 1-2*(x*x + y*y))
            pitch = torch.asin(torch.clip(2*(w*y - z*x), -1, 1))
            yaw = torch.atan2(2*(w*z + x*y), 1-2*(y*y + z*z))
            return roll, pitch, yaw
        
        roll, pitch, yaw = euler_from_quaternion(root_rot)
        root_vel = quat_rotate_inverse(root_rot, root_vel)
        root_ang_vel = quat_rotate_inverse(root_rot, root_ang_vel)
        
        # 重塑观测
        root_pos = root_pos.reshape(1, num_steps, 3)
        root_vel = root_vel.reshape(1, num_steps, 3) 
        root_ang_vel = root_ang_vel.reshape(1, num_steps, 3)
        dof_pos = dof_pos.reshape(1, num_steps, -1)
        roll = roll.reshape(1, num_steps, 1)
        pitch = pitch.reshape(1, num_steps, 1)
        
        mimic_obs = torch.cat((root_pos[...,2:3], roll, pitch, root_vel, 
                              root_ang_vel[...,2:3], dof_pos), dim=-1)
        
        return mimic_obs.reshape(1, -1).detach().cpu().numpy().squeeze()
    
    def _quat_to_euler(self, quat):
        """四元数转欧拉角"""
        x, y, z, w = quat
        roll = np.arctan2(2*(w*x + y*z), 1-2*(x*x + y*y))
        pitch = np.arcsin(np.clip(2*(w*y - z*x), -1, 1))
        yaw = np.arctan2(2*(w*z + x*y), 1-2*(y*y + z*z))
        return np.array([roll, pitch, yaw])
    
    def _send_commands(self, target_positions, current_state):
        """
        机器人控制命令发送 - 区分真实与模拟控制
        """
        if not UNITREE_AVAILABLE:
            return self._send_simulation_commands(target_positions, current_state)
        else:
            return self._send_real_robot_commands(target_positions, current_state)
    
    def _send_simulation_commands(self, target_positions, current_state):
        """
        模拟控制命令发送 - 用于测试和开发
        
        【控制目标】: 仿真测试模式 - 无实际硬件控制
        【SDK调用】: 无 (纯Python日志记录)
        【用途】: 算法验证、离线测试、开发调试
        """
        logging.debug("🎮 控制模式: 仿真模拟控制")
        
        # 【模拟控制】仅记录控制指令，不发送到硬件
        control_data = {
            'target_positions': target_positions.tolist(),
            'timestamp': time.time(),
            'control_mode': 'simulation',
            'control_target': 'simulation_mock',
            'sdk_calls': 'none (python logging only)'
        }
        
        # 【日志记录】控制指令验证 (替代真实硬件控制)
        if self.curr_timestep % 100 == 0:  # 每2秒记录一次
            logging.debug(f"📝 模拟控制指令: 关节目标位置 = {target_positions[:6].tolist():.3f}...")
        
        return control_data
    
    def _send_real_robot_commands(self, target_positions, current_state):
        """
        真实机器人控制命令发送 - 发送到G1机器人硬件
        
        【控制目标】: Unitree G1机器人23个电机硬件
        【SDK调用】: unitree_sdk2_python DDS通信发送
        【硬件详情】:
        - 23个电机: 伺服电机 + 减速器 + 电机驱动器
        - DDS通信: 实时数据分发服务协议
        - PD控制器: 位置-速度双环控制算法
        """
        logging.debug("🤖 控制模式: 真实机器人硬件")
        
        # 【数据来源】G1机器人实时状态数据
        current_pos = current_state['joint_pos']  # 来自电机编码器
        current_vel = current_state['joint_vel']  # 来自电机速度反馈
        
        # 生成23个电机的PD控制命令
        control_torques = []
        for i in range(23):
            # 【算法】经典PD控制器计算
            pos_error = target_positions[i] - current_pos[i]  # 位置误差
            vel_error = 0.0 - current_vel[i]                  # 速度误差(目标速度为0)
            tau = pos_error * self.stiffness[i] + vel_error * self.damping[i]  # PD计算
            
            # 【安全限制】基于G1机器人硬件扭矩限制
            tau = np.clip(tau, -self.torque_limits[i], self.torque_limits[i])
            control_torques.append(tau)
            
            # 【SDK调用】unitree_sdk2_python -> 填充电机命令结构
            # 【硬件目标】G1机器人第{i}号伺服电机驱动器
            self.low_cmd.motor_cmd[i].q = target_positions[i]    # 目标位置
            self.low_cmd.motor_cmd[i].dq = 0.0                   # 目标速度
            self.low_cmd.motor_cmd[i].tau = tau                  # 控制扭矩
        
        # 【SDK调用】unitree_sdk2_python.utils.crc.CRC.Crc32 - 计算校验和
        # 【用途】确保DDS消息完整性，防止传输错误
        self.low_cmd.crc = CRC.Crc32(self.low_cmd, len(self.low_cmd) - 4)
        
        # 【SDK调用】unitree_sdk2_python.ChannelPublisher.write() - DDS发送命令
        # 【通信协议】DDS (Data Distribution Service) 实时通信
        # 【硬件目标】G1机器人23个电机硬件控制器
        # 【物理效果】伺服电机根据PD控制指令执行位置控制
        self.cmd_pub.write(self.low_cmd)
        
        return {
            'target_positions': target_positions.tolist(),
            'control_torques': control_torques,
            'timestamp': time.time(),
            'control_mode': 'real_robot',
            'control_target': 'unitree_g1_motors',
            'sdk_calls': 'unitree_sdk2_python DDS communication'
        }
    
    def _safety_check(self, target_pos, current_pos, dt):
        """安全检查"""
        # 位置限制
        pos_limits = np.array([
            [-1.57,0.5], [-0.5,0.5], [-1.57,1.57], [0.0,2.5], [-0.8,0.8], [-0.5,0.5],  # 左腿
            [-1.57,0.5], [-0.5,0.5], [-1.57,1.57], [0.0,2.5], [-0.8,0.8], [-0.5,0.5],  # 右腿
            [-1.57,1.57], [-0.5,0.5], [-0.5,0.5],                                        # 躯干
            [-3.14,3.14], [-1.57,1.57], [-1.57,1.57], [0.0,2.6],                       # 左臂
            [-3.14,3.14], [-1.57,1.57], [-1.57,1.57], [0.0,2.6],                       # 右臂
        ])
        
        for i, (pos, (low, high)) in enumerate(zip(target_pos, pos_limits)):
            if not (low <= pos <= high):
                raise Exception(f"Joint {i} position {pos:.3f} out of limits [{low:.3f}, {high:.3f}]")
        
        # 速度限制
        vel = (target_pos - current_pos) / dt
        max_vel = 10.0
        for i, v in enumerate(vel):
            if abs(v) > max_vel:
                raise Exception(f"Joint {i} velocity {v:.3f} exceeds limit {max_vel}")
        
        return True
    
    def control_loop(self):
        """主控制循环"""
        logging.info("Starting control loop...")
        self.running = True
        
        while self.running:
            start_time = time.time()
            
            try:
                # === 1. 读取状态 ===
                # 【SDK调用】unitree_sdk2_python.ChannelSubscriber.read() 或 仿真模式
                # 【数据来源】G1机器人硬件传感器 或 仿真零值
                robot_state = self._read_robot_state()
                if robot_state is None:
                    time.sleep(0.001)
                    continue
                
                # === 2. 获取运动参考 ===
                # 【SDK调用】自定义motion_lib.calc_motion_frame()
                # 【数据来源】PKL文件中的预录制人类动作数据
                mimic_obs = self._get_mimic_obs(self.curr_timestep)
                
                # === 3. 构建观测 ===
                # 【数据来源】G1机器人IMU传感器 + 关节编码器 + 历史缓存
                rpy = self._quat_to_euler(robot_state['imu_quat'])  # IMU四元数转欧拉角
                obs_dof_vel = robot_state['joint_vel'].copy()       # 关节速度数据
                obs_dof_vel[[4,5,10,11]] = 0.0  # 清零脚踝扭矩传感器(提高鲁棒性)
                
                obs_prop = np.concatenate([
                    robot_state['imu_gyro'] * 0.25,
                    rpy[:2],
                    (robot_state['joint_pos'] - self.default_dof_pos) * 1.0,
                    obs_dof_vel * 0.05,
                    self.last_action,
                ])
                
                obs_hist = np.array(self.proprio_history_buf).flatten()
                obs_buf = np.concatenate([mimic_obs, obs_prop, obs_hist])
                
                # === 4. 策略推理 ===
                # 【SDK调用】PyTorch JIT模型前向传播
                # 【数据来源】在MuJoCo仿真中训练的策略网络权重
                obs_tensor = torch.from_numpy(obs_buf).float().unsqueeze(0).to(self.device)
                with torch.no_grad():
                    raw_action = self.policy_jit(obs_tensor).cpu().numpy().squeeze()
                
                # 5. 动作处理
                self.last_action = raw_action.copy()
                raw_action = np.clip(raw_action, -10., 10.)
                scaled_actions = raw_action * self.action_scale
                target_positions = scaled_actions + self.default_dof_pos
                
                # 6. 安全检查
                self._safety_check(target_positions, robot_state['joint_pos'], self.control_dt)
                
                # 7. 发送命令 - 区分真实控制与模拟控制
                # 【SDK调用】unitree_sdk2_python DDS通信 或 仿真日志记录
                # 【控制目标】G1机器人23个电机硬件 或 仿真测试环境
                control_result = self._send_commands(target_positions, robot_state)
                
                # 8. 更新历史
                self.proprio_history_buf.append(obs_prop)
                self.curr_timestep += 1
                
                # 9. 数据记录
                if self.curr_timestep % 50 == 0:  # 每秒记录
                    step_data = {
                        'timestep': self.curr_timestep,
                        'joint_pos': robot_state['joint_pos'].tolist(),
                        'target_pos': target_positions.tolist(),
                        'action': raw_action.tolist(),
                        'imu_quat': robot_state['imu_quat'].tolist(),
                    }
                    self.data_log.append(step_data)
                    
                    freq = 1.0 / (time.time() - start_time + 1e-6)
                    logging.info(f"Step: {self.curr_timestep}, Freq: {freq:.1f}Hz")
                
                # 10. 控制频率
                elapsed = time.time() - start_time
                sleep_time = max(0, self.control_dt - elapsed)
                time.sleep(sleep_time)
                
            except KeyboardInterrupt:
                logging.info("Received keyboard interrupt")
                break
            except Exception as e:
                logging.error(f"Control error: {e}")
                self.emergency_stop()
                break
        
        self.stop()
    
    def emergency_stop(self):
        """
        紧急停止 - 区分真实与模拟控制
        """
        if not UNITREE_AVAILABLE:
            return self._emergency_stop_simulation()
        else:
            return self._emergency_stop_real_robot()
    
    def _emergency_stop_simulation(self):
        """
        模拟模式紧急停止
        
        【控制目标】: 仿真测试模式 - 停止模拟控制循环
        【SDK调用】: 无 (纯Python状态控制)
        【用途】: 开发测试时的安全退出
        """
        logging.warning("🚨 仿真模式紧急停止执行")
        self.running = False
        
        emergency_data = {
            'timestamp': time.time(),
            'stop_mode': 'simulation',
            'stop_target': 'simulation_mock',
            'sdk_calls': 'none (python state control only)'
        }
        
        return emergency_data
    
    def _emergency_stop_real_robot(self):
        """
        真实机器人紧急停止
        
        【控制目标】: Unitree G1机器人23个电机硬件
        【SDK调用】: unitree_sdk2_python DDS通信
        【硬件效果】: 立即切换电机为零扭矩模式，停止所有运动
        """
        logging.warning("🚨 真实机器人紧急停止执行")
        
        # 【SDK调用】unitree_sdk2_python -> 设置所有电机为被动模式
        # 【硬件目标】G1机器人23个伺服电机驱动器
        # 【物理效果】电机进入零扭矩状态，机器人自然下垂
        for i in range(23):
            self.low_cmd.motor_cmd[i].mode = 0x00  # 被动模式 (零扭矩)
            self.low_cmd.motor_cmd[i].tau = 0.0    # 扭矩清零
        
        # 【SDK调用】unitree_sdk2_python.utils.crc.CRC.Crc32
        self.low_cmd.crc = CRC.Crc32(self.low_cmd, len(self.low_cmd) - 4)
        
        # 【SDK调用】unitree_sdk2_python.ChannelPublisher.write()
        # 【通信协议】DDS实时通信发送紧急停止命令
        # 【硬件目标】G1机器人23个电机控制器
        self.cmd_pub.write(self.low_cmd)
        
        self.running = False
        
        return {
            'timestamp': time.time(),
            'stop_mode': 'real_robot', 
            'stop_target': 'unitree_g1_motors',
            'sdk_calls': 'unitree_sdk2_python DDS emergency stop'
        }
    
    def stop(self):
        """停止控制器"""
        self.running = False
        
        # 保存数据日志
        if self.data_log:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            log_file = f"logs/control_data_{timestamp}.json"
            os.makedirs("logs", exist_ok=True)
            with open(log_file, 'w') as f:
                json.dump(self.data_log, f, indent=2)
            logging.info(f"Data saved to {log_file}")
        
        logging.info("Controller stopped")

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description='G1 Sim2Real Unified Controller')
    parser.add_argument('--motion', type=str, default='walk_stand.pkl', help='Motion file')
    parser.add_argument('--interface', type=str, default='lo', help='Network interface (lo/enp3s0)')
    parser.add_argument('--domain_id', type=int, default=1, help='DDS domain ID (1=sim, 0=real)')
    parser.add_argument('--device', type=str, default='cpu', help='Device (cpu/cuda)')
    
    args = parser.parse_args()
    
    # 文件路径
    policy_path = "assets/pretrained_checkpoints/pretrained.pt"
    motion_path = f"assets/motions/{args.motion}"
    
    if not os.path.exists(policy_path):
        print(f"❌ Policy file not found: {policy_path}")
        return
    
    if not os.path.exists(motion_path):
        print(f"❌ Motion file not found: {motion_path}")
        return
    
    # 创建控制器
    controller = G1Sim2RealController(
        policy_path=policy_path,
        motion_path=motion_path,
        interface=args.interface,
        domain_id=args.domain_id,
        device=args.device
    )
    
    print(f"🚀 Starting controller...")
    print(f"   Motion: {args.motion}")
    print(f"   Interface: {args.interface}")
    print(f"   Mode: {'Real Robot' if args.interface != 'lo' else 'Simulation'}")
    print(f"   Press Ctrl+C to stop")
    
    try:
        controller.control_loop()
    except KeyboardInterrupt:
        print("\nStopping controller...")
        controller.stop()

if __name__ == "__main__":
    main()